import numpy as np
import os
import torch
import cv2

import matplotlib.pyplot as plt


def load_data():
    pass


def imgBrightness(img1, c, b):
    rows, cols, channels = img1.shape
    blank = np.zeros([rows, cols, channels], img1.dtype)
    rst = cv2.addWeighted(img1, c, blank, 1 - c, b)
    return rst


if __name__ == '__main__':

    path = '/data/fg2021/DG-13-x5'
    x5_path = '/data/fg2021/DG-13-x5-depth'
    folders = os.listdir(path)

    for folder in folders:
        folders_path = path + '/' + folder
        files = os.listdir(folders_path)
        x5_path_folder = x5_path + '/' + folder

        try:
            os.mkdir(x5_path_folder)

        except:
            pass

        for file in files:
            name = folders_path + '/' + file
            data = np.load(name)
            data = data.transpose(1, 2, 3, 0)
            seq = []

            for frame in data:
                depth = np.zeros((256, 256, 1), dtype=np.uint8)
                depth[:, :, 0] = frame[:, :, 3]
                seq.append(depth)
            seq_np = np.array(seq).transpose(3, 0, 1, 2)

            # seq_rgb_np = np.array(seq_rgb).transpose(3, 0, 1, 2)
            # print(file)
            # print(seq_rgb_np.shape)
            # cv2.waitKey(-1)

            np.save(x5_path_folder + '/' + file, seq_np)

        print(folder + " finish")
